import  numpy as np
import random
from sklearn.model_selection import KFold
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestClassifier as RFC
from imblearn.under_sampling import RandomUnderSampler
from xgboost import XGBClassifier as XGBC
from sklearn.metrics import log_loss
import pandas as pd
import fig
import os

# y true, y predicted
def accuracy(yp, yt):
    cor = 0
    for i in range(len(yp)):
        if yp[i] == yt[i]:
            cor=cor+1
    return cor/len(yp)

# 读取数据
def load_data(path):
    data = pd.read_csv(path, delimiter=r",")
    values = data.values[:, 0:1].ravel()
    valid = False
    X = []
    for data_idx in range(len(values)):
        alloy, alpha_st, num_st = [], -1, -1
        value = values[data_idx]
        leng = len(value)
        for i in range(leng):
            if value[i].isalpha() :
                if alpha_st == -1:
                    alpha_st = i
                if num_st != -1 :
                    a = value[num_st:i]
                    alloy.append(float(a))
                    num_st = -1
                if i == leng-1:
                    alloy.append(value[num_st:leng])
            else:
                if num_st == -1:
                    num_st = i
                if alpha_st != -1:
                    a = value[alpha_st:i]
                    alloy.append(a)
                    alpha_st = -1
        alloy.append(float(value[num_st:leng]))

        node = np.zeros((100, 1))
        used_idx = []

        for i in range(0, len(alloy), 2):
            idx = elements.get(alloy[i], len(elements))
            if idx == len(elements):
                elements[alloy[i]] = idx
            used_idx.append(idx)
            node[idx] = alloy[i+1]
        X.append(node.ravel())
    return np.array(X), np.array(data.values[:, -1].ravel())

def generate_dataset():
    X, Y = np.array([]), np.array([])
    file_list = os.listdir(DATA_PATH)
    file_list.sort()
    for file in file_list:
        if file not in ["test_alloy_idx.csv", "test_alloy_sys_idx.csv"]:
            x_new, y_new = load_data(DATA_PATH + file)
            X = x_new if len(X)==0 else np.vstack((X, x_new))
            Y = y_new if len(Y)==0 else np.hstack((Y, y_new))
    x_new, y_new = load_data(DATA_PATH + "1877-am.csv")
    for i in range(len(y_new)):
        y_new[i]=1
    X = np.vstack((X, x_new))
    Y = np.hstack((Y, y_new))
    Y = Y.astype('int')
    rus = RandomUnderSampler(random_state=0, replacement=True)
    x_Usampled, y_Usampled = rus.fit_resample(X,Y)

    elen = len(elements)
    scaler = MinMaxScaler(feature_range=(0, 1))
    x_Usampled = scaler.fit_transform(x_Usampled)[:, 0:elen]
    return x_Usampled, y_Usampled

def split_index(n_total, ratio, shuffle=True):
    idx_list = list(range(0, n_total))
    offset = int(n_total * ratio)
    random.shuffle(idx_list)
    sublist1 = idx_list[:offset]
    sublist2 = idx_list[offset:]
    return sublist1, sublist2

def get_alloy_systems(X):
    alloy_systems = []
    for a in X:
        idx = [i for i, e in enumerate(a) if e != 0]
        if idx not in alloy_systems:
            alloy_systems.append(idx)
    return alloy_systems

def split_by_alloy_system_index(X, Y, alloy_system_idx):
    v_alloy_idx = []
    for i in range(len(X)):
        alloy_item = X[i]
        for alloy_sys_item in alloy_system_idx:
            if np.count_nonzero(alloy_item) == len(alloy_sys_item):
                tag = True
                for e in alloy_sys_item:
                    if alloy_item[e] == 0:
                        tag = False
                        break
                if tag:
                    v_alloy_idx.append(i)
    return v_alloy_idx

#输入划分多少，以及是否按照alloy system划分，得到test index
def split_test_data(X, Y, split_rate):
    alloy_systems = get_alloy_systems(X)
    v_sys_idx, t_sys_idx = split_index(len(alloy_systems), split_rate, shuffle=True)
    valid_alloy_system_idx = []
    for idx in v_sys_idx:
        valid_alloy_system_idx.append(alloy_systems[idx])
    test_alloy_sys_index = split_by_alloy_system_index(X, Y, valid_alloy_system_idx)
    out_data = pd.DataFrame(data=test_alloy_sys_index, columns=["idx"])
    out_data.to_csv(DATA_PATH + "test_alloy_sys_idx.csv", index=False)

    test_alloy_index, train_alloy_index = split_index(len(Y), split_rate, shuffle=True)
    out_data = pd.DataFrame(data=test_alloy_index, columns=["idx"])
    out_data.to_csv(DATA_PATH + "test_alloy_idx.csv", index=False)

def kfold_alloy_system_valid(X, Y, model):
    alloy_systems = get_alloy_systems(X)

    ACC, loss, n_splits = 0, 0, 10
    kf = KFold(n_splits, shuffle=True)
    for train_index, valid_index in kf.split(alloy_systems):
        valid_alloy_system_idx = []
        for idx in valid_index:
            valid_alloy_system_idx.append(alloy_systems[idx])

        v_alloy_sys_index = split_by_alloy_system_index(X, Y, valid_alloy_system_idx)
        x_valid, y_valid = np.array(X[v_alloy_sys_index]), np.array(Y[v_alloy_sys_index])

        x_train, y_train = [], []
        for i in range(len(Y)):
            if i not in v_alloy_sys_index:
                x_train.append(X[i])
                y_train.append(Y[i])

        model.fit(x_train, y_train)  # 训练
        y_pred = model.predict(x_valid)
        acc = accuracy(y_pred, y_valid)
        #print(acc)
        ACC = ACC + acc
        loss = loss + log_loss(y_valid, model.predict_proba(x_valid))
    print(f'kfold alloy systems acc {ACC/n_splits} loss {loss/n_splits}')


def kfold_alloy_valid(X, Y, model):
    ACC, loss, n_splits = 0, 0, 10
    kf = KFold(n_splits, shuffle=True)
    for train_index, valid_index in kf.split(X):
        x_train, y_train = X[train_index], Y[train_index]
        x_valid, y_valid = X[valid_index], Y[valid_index]

        x_train = x_train.astype('float')
        y_train = y_train.astype('float')
        x_valid = x_valid.astype('float')
        y_valid = y_valid.astype('float')

        model.fit(x_train, y_train)  # 训练
        y_pred = model.predict(x_valid)
        y_pred_prob = model.predict_proba(x_valid)

        acc = accuracy(y_pred, y_valid)
        #print(acc)
        ACC = ACC + acc
        loss = loss + log_loss(y_valid, y_pred_prob)
    print(f'kfold acc {ACC/n_splits} {loss/n_splits}')


rfc = RFC(n_estimators=50)
xgbc = XGBC(n_estimators=350,
            max_depth=10,
            min_child_weight=1,
            subsample=1,
            learning_rate=0.1,
            )

def set_elements():
    data = pd.read_csv("data/elements.csv", delimiter=r",")
    A = data.values[:, 0:1].ravel()
    B = data.values[:, -1].ravel()
    for i in range(len(A)):
        elements[A[i]] = B[i]

def predict(test_x):
    #elements 要对应起来
    set_elements()
    x_train, y_train = generate_dataset()
    x_train = x_train.astype('float')
    y_train = y_train.astype('float')
    xgbc.fit(x_train, y_train)  # 训练

    print(xgbc)
    y_pred = xgbc.predict(test_x)
    num_am = 0
    for i in range(len(y_pred)):
        if y_pred[i] == 1:
            num_am = num_am + 1

    print(f'am num {num_am} {num_am/len(test_x)}')
    return y_pred

DATA_PATH = "data/xgboost-am/"
elements = dict()

#分离出来两种test文件，输出test index
NEW_TEST_DATA = False
TEST = True
def main():
    #load elements idx
    set_elements()

    X, Y = generate_dataset()
    for key in elements.keys():
        print(f"{key}, {elements[key]}")

    if NEW_TEST_DATA:
         split_test_data(X, Y, 0.1)

    #data loading
    test_aly_idx = pd.read_csv(DATA_PATH + "test_alloy_idx.csv", delimiter=r",").values
    test_aly_sys_idx = pd.read_csv(DATA_PATH + "test_alloy_sys_idx.csv", delimiter=r",").values

    w = len(elements)
    X_alloy_test, Y_alloy_test = X[test_aly_idx].reshape(len(test_aly_idx),w), Y[test_aly_idx]
    X_alloy_sys_test = np.array(X[test_aly_sys_idx]).reshape(len(test_aly_sys_idx), w)
    Y_alloy_sys_test = np.array(Y[test_aly_sys_idx])
    X_tv, Y_tv = [], [] # train and validation
    for i in range(len(Y)):
        if (i not in test_aly_idx) and (i not in test_aly_sys_idx):
            X_tv.append(X[i])
            Y_tv.append(Y[i])

    X_tv, Y_tv = np.array(X_tv), np.array(Y_tv)
    print(f"test alloy system {len(X_alloy_sys_test)} test alloy {len(Y_alloy_test)}")

    #valid alloy
    kfold_alloy_valid(X_tv, Y_tv, xgbc)
    # valid alloy system
    kfold_alloy_system_valid(X_tv, Y_tv, xgbc)

    if TEST:
        # test alloy
        xgbc.fit(X_tv, Y_tv)
        y_pred_alloy = xgbc.predict(X_alloy_test)
        y_pred_alloy_proba = xgbc.predict_proba(X_alloy_test)
        print(f'data test acc {accuracy(Y_alloy_test, y_pred_alloy)} '
              f'loss {log_loss(Y_alloy_test, y_pred_alloy_proba)}')

        # test alloy system
        xgbc.fit(X_tv, Y_tv)
        y_pred_aly_sys = xgbc.predict(X_alloy_sys_test)
        y_pred_alloy_sys_proba = xgbc.predict_proba(X_alloy_sys_test)
        print(f'alloy system test acc {accuracy(Y_alloy_sys_test, y_pred_aly_sys)} '
              f'loss {log_loss(Y_alloy_sys_test, y_pred_alloy_sys_proba)}')

if __name__ == '__main__':
    main()